import torch
import torchvision
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from model import network

cuda = True
device = torch.device("cuda" if cuda else "cpu")


class OriPlace(nn.Module):
    def __init__(self, backbone):
        super(OriPlace, self).__init__()
        self.sat_backbone, sat_features_dim = network.get_backbone(backbone)
        self.grd_backbone, grd_features_dim = network.get_backbone(backbone)
        self.sdm = network.SDM(128, sat_features_dim)

    def forward(self, sat, grd):
        sat_coding = self.sat_backbone(sat)
        grd_coding = self.grd_backbone(grd)
        sat_coding = F.normalize(sat_coding, dim=1)
        grd_coding = F.normalize(grd_coding, dim=1)
        ori, option_ori = self.sdm(grd_coding, sat_coding)

        return ori, option_ori


class CosPlace(nn.Module):
    def __init__(self, backbone, dim):
        super(CosPlace, self).__init__()
        self.sat_geolocnet = network.GeoLocalizationNet(backbone, dim)
        self.grd_geolocnet = network.GeoLocalizationNet(backbone, dim)

    def forward(self, x_sat, x_grd):
        sat_code = self.sat_geolocnet.forward(x_sat)
        grd_code = self.grd_geolocnet.forward(x_grd)
        return sat_code, grd_code
